#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Analyze zone results and visualize
===================================
- 读取推理结果 JSON
- 根据 path2label.json 添加 global label
- 可视化原图 + 裁剪框
- 保存每个 case 的裁剪图 + info.json
"""

import os
import json
import argparse
from PIL import Image, ImageDraw, ImageFont


def analyze_zone(input_json, path2label_json, output_dir, max_num, font_path="./SarasaMonoCL-Regular.ttf"):
    # 加载 path2label
    if not os.path.exists(path2label_json):
        raise FileNotFoundError(f"path2label 文件不存在: {path2label_json}")
    with open(path2label_json, "r", encoding="utf-8") as f:
        path2label = json.load(f)

    # 加载输入结果
    if not os.path.exists(input_json):
        raise FileNotFoundError(f"输入文件不存在: {input_json}")
    with open(input_json, "r", encoding="utf-8") as f:
        data = json.load(f)

    # 字体
    try:
        font = ImageFont.truetype(font_path, 20)
    except Exception as e:
        raise RuntimeError(f"字体加载失败: {font_path}, 错误: {e}")

    counters = {"accept": 0, "reject": 0, "uncertain": 0, "failed": 0}
    os.makedirs(output_dir, exist_ok=True)

    for item in data:
        image_path = item["image_path"]
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"原图不存在: {image_path}")

        global_label = path2label.get(image_path, "unknown")
        img = Image.open(image_path).convert("RGB")

        # --- 判定该图类别 ---
        cat = "accept"
        has_reject, has_uncertain, has_failed = False, False, False

        for bbox in item["bboxes"]:
            if "result" not in bbox or "is_correct" not in bbox["result"]:
                raise KeyError(f"[ERROR] bbox 缺少 result/is_correct 字段: {bbox}")
            ic = str(bbox["result"]["is_correct"]).lower()
            if ic == "failed":
                has_failed = True
            elif ic == "uncertain":
                has_uncertain = True
            elif ic == "false":
                has_reject = True

        if has_failed:
            cat = "failed"
        elif has_uncertain:
            cat = "uncertain"
        elif has_reject:
            cat = "reject"
        else:
            cat = "accept"

        # --- 判断是否保留 ---
        keep_image = False
        if cat == "failed":
            keep_image = True  # 强制保留
        else:
            if max_num == -1 or counters[cat] < max_num:
                keep_image = True

        if not keep_image:
            continue

        # --- 可视化与保存 ---
        case_name = os.path.splitext(os.path.basename(image_path))[0]
        case_dir = os.path.join(output_dir,f"{cat}_{case_name}")
        os.makedirs(case_dir, exist_ok=True)

        draw = ImageDraw.Draw(img)
        info = {"image_path": image_path, "global_label": global_label, "category": cat, "bboxes": []}

        for idx, bbox in enumerate(item["bboxes"]):
            result = bbox["result"]
            xyxy = bbox["xyxy"]

            draw.rectangle(xyxy, outline="red", width=3)
            draw.text((xyxy[0], xyxy[1] - 25), str(idx), font=font, fill="red")

            crop = img.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
            crop.save(os.path.join(case_dir, f"{idx}.jpg"), quality=95)

            info["bboxes"].append({
                "id": idx,
                "xyxy": xyxy,
                "description": bbox["description"],
                "is_correct": result["is_correct"],
                "confidence": result["confidence"],
                "region_description": result["region_description"],
                "rationale": result["rationale"]
            })

        counters[cat] += 1
        img.save(os.path.join(case_dir, f"original_{case_name}.jpg"), quality=95)
        with open(os.path.join(case_dir, "info.json"), "w", encoding="utf-8") as f:
            json.dump(info, f, ensure_ascii=False, indent=2)

    print(f"[INFO] 可视化完成，结果保存在 {output_dir}")
    print(f"[INFO] per-class 数量统计 (按 image 加权): {counters}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Analyze and visualize zone results into JSON.")
    parser.add_argument("--input_json", type=str, default="./experiments/output.json", help="推理结果 JSON 文件")
    parser.add_argument("--path2label_json", type=str, default="./experiments/path2label.json", help="全局 path2label.json")
    parser.add_argument("--output_dir", type=str, default="./experiments/analyze_zone", help="输出目录")
    parser.add_argument("--max_num", type=int, default=50, help="每个类别最多可视化多少张, -1 表示全部")
    parser.add_argument("--font_path", type=str, default="./SarasaMonoCL-Regular.ttf", help="字体路径")

    args = parser.parse_args()
    analyze_zone(args.input_json, args.path2label_json, args.output_dir, args.max_num, args.font_path)
